#!/usr/bin/env python3
"""
LLM Prompt Risk Scanner Setup Script
Developed by Syed Ali Turab

This script helps users set up the LLM Prompt Risk Scanner with interactive API key configuration.
"""

import os
import sys
import subprocess
from pathlib import Path

def print_banner():
    """Print the setup banner"""
    print("""
    🛡️  LLM PROMPT RISK SCANNER SETUP
    =====================================
    Multi-Provider Security Analysis Tool
    Developed by Syed Ali Turab
    =====================================
    """)

def check_python_version():
    """Check if Python version is compatible"""
    if sys.version_info < (3, 8):
        print("❌ Error: Python 3.8 or higher is required.")
        print(f"   Current version: {sys.version}")
        sys.exit(1)
    print(f"✅ Python version: {sys.version.split()[0]}")

def install_requirements():
    """Install required packages"""
    print("\n📦 Installing required packages...")
    try:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-r", "requirements.txt"])
        print("✅ All packages installed successfully!")
    except subprocess.CalledProcessError:
        print("❌ Failed to install packages. Please run: pip install -r requirements.txt")
        return False
    return True

def setup_env_file():
    """Interactive setup of environment variables"""
    print("\n🔑 API Key Configuration")
    print("=" * 40)
    
    env_file = Path(".env")
    env_content = []
    
    # API providers configuration
    providers = {
        "TOGETHER_API_KEY": {
            "name": "Together AI",
            "url": "https://api.together.xyz",
            "description": "Recommended - Free tier with 5M tokens/month",
            "required": False
        },
        "OPENAI_API_KEY": {
            "name": "OpenAI",
            "url": "https://platform.openai.com/api-keys",
            "description": "Popular choice - Requires billing setup",
            "required": False
        },
        "ANTHROPIC_API_KEY": {
            "name": "Anthropic Claude",
            "url": "https://console.anthropic.com",
            "description": "High quality - $5 free credit",
            "required": False
        },
        "DEEPSEEK_API_KEY": {
            "name": "DeepSeek",
            "url": "https://platform.deepseek.com",
            "description": "Cost-effective - Competitive pricing",
            "required": False
        }
    }
    
    print("You need at least ONE API key to use the scanner.")
    print("Together AI is recommended for beginners (free tier).\n")
    
    configured_count = 0
    
    for key, info in providers.items():
        print(f"🔧 {info['name']}")
        print(f"   Description: {info['description']}")
        print(f"   Get key from: {info['url']}")
        
        while True:
            api_key = input(f"   Enter {info['name']} API key (or press Enter to skip): ").strip()
            
            if not api_key:
                print("   ⏭️  Skipped")
                break
            
            if len(api_key) < 10:  # Basic validation
                print("   ❌ API key seems too short. Please check and try again.")
                continue
            
            env_content.append(f"{key}={api_key}")
            print("   ✅ Added successfully!")
            configured_count += 1
            break
        
        print()
    
    if configured_count == 0:
        print("❌ No API keys configured. You'll need at least one to use the scanner.")
        print("   You can always add them later to the .env file.")
        return False
    
    # Write .env file
    try:
        with open(env_file, "w") as f:
            f.write("# LLM Prompt Risk Scanner - API Keys\n")
            f.write("# Generated by setup script\n\n")
            for line in env_content:
                f.write(line + "\n")
        
        print(f"✅ Environment file created with {configured_count} API key(s)")
        return True
    
    except Exception as e:
        print(f"❌ Failed to create .env file: {e}")
        return False

def run_test():
    """Test the installation"""
    print("\n🧪 Testing installation...")
    
    try:
        # Test imports
        import streamlit
        import requests
        import dotenv
        print("✅ Core packages imported successfully")
        
        # Test optional imports
        try:
            import openai
            print("✅ OpenAI package available")
        except ImportError:
            print("⚠️  OpenAI package not available (install with: pip install openai)")
        
        try:
            import anthropic
            print("✅ Anthropic package available")
        except ImportError:
            print("⚠️  Anthropic package not available (install with: pip install anthropic)")
        
        try:
            import together
            print("✅ Together package available")
        except ImportError:
            print("⚠️  Together package not available (install with: pip install together)")
        
        return True
    
    except ImportError as e:
        print(f"❌ Import error: {e}")
        return False

def main():
    """Main setup function"""
    print_banner()
    
    # Check Python version
    check_python_version()
    
    # Install requirements
    if not install_requirements():
        sys.exit(1)
    
    # Setup environment file
    env_setup = setup_env_file()
    
    # Run tests
    test_passed = run_test()
    
    # Final instructions
    print("\n" + "=" * 50)
    print("🎉 SETUP COMPLETE!")
    print("=" * 50)
    
    if env_setup and test_passed:
        print("✅ Everything is ready to go!")
        print("\n🚀 To start the scanner, run:")
        print("   streamlit run app.py")
        print("\n📖 The app will open at: http://localhost:8501")
    else:
        print("⚠️  Setup completed with some issues.")
        print("   Please check the messages above and resolve any problems.")
    
    print("\n📚 Additional Resources:")
    print("   • README.md - Complete documentation")
    print("   • .env.example - Environment variable examples")
    print("   • GitHub Issues - Report problems or get help")
    
    print("\n🛡️ Happy scanning!")

if __name__ == "__main__":
    main()